#include "data_gen.h"
  
// Generate a random message symbol.  Symbol must lie in the field GF(2^m)
int gen_symbol(int m) {

    int sym;
    int i;
    int r;

    sym = 0;
    for (i = 0; i < m; i++) {
        r = (rand()) / (RAND_MAX / ((1<<m)-2));
//        r = 1 & (r >> 11); // Need to grab a random bit from somewhere in the
                           // middle of the number because low bits of rand() 
                           // have poor randomness.
 //       sym = r * (1 << i);
        sym = r;
    }

    return sym;

}

// Generate a data word and encode into a code word
Vector * gen_codeword(int m, int t, Matrix *add_one, Vector *msg, Vector *g) {

    int n;
    Vector *T; // Code word to be transmitted
    Vector *M; // Message word
    int i;

    n = (1 << m) - 1;

    // Create a random message word
    M = allocate_vector(n-2*t);
    for (i = 0; i < M->len; i++) {
        M->data[i] = gen_symbol(m);   
    }

    // Create a code word using the random message word
    T = RS_encoder(m, add_one, t, M, g); 
    
    free(M);

    return T;

}

// Add errors to a message polynomial
Vector * add_error(double e_prob, int e_per_word, double *e_pmf, int m, 
                                                        int mode, Vector *T) {
   
    Vector *R; // "Received" polynomial (polynomial with errors).
    int i,j;
    double k;
    int *err_locs;

    R = copy_vector(T);

    if (mode == FIXED_ERR_PROB) {
        for (i = 0; i < R->len; i++) {
            k = rand();
            if (k > (1-e_prob) * RAND_MAX) { // Error occurs
                
                // Create random error
                R->data[i] = (R->data[i] + (rand() % ((1<<m)-1))) % ((1<<m)-1);
                
            }
        }
    }
    else { // mode == PROB_PER_SYMBOL
        // Allocate an array to house the error locations
        err_locs = (int*)malloc((T->len)*sizeof(int));

        // Generate the errors
        for (i = 0; i < e_per_word; i++) {

            // Randomly choose the error location according to e_pmf            
            k = rand() / RAND_MAX;
            for (j = 0; j < T->len; j++) {
                if (k <= e_pmf[j]) {

                    // Create random error at this location
                    R->data[i] = (R->data[i] + (rand() % ((1<<m)-1))) % ((1<<m)-1);
                    break;
                }
            }
        }
    }

    return R;
}

// Create a data file of length data_length.
int create_data(int     m, 
                Matrix  *add_one, 
                int     t, 
                int     data_length,
                double  e_prob,
                int     e_per_word,
                double  *e_pmf,
                int     mode, 
                char    *fname
                ) {
    
    FILE *file;
    int i,j,k;
    Vector * codeword;  // Codeword polynomial
    Vector * R;         // "Received" polynomial (poly + errors)
    Vector * g;         // Generator polynomial
    Vector * M;         // Message polynomial

    // Seed the random number generator
    srand(time(0));

    // Derive the generator polynomial
    g = RS_generator(m, add_one, t);

    // Open the output file for writing
    file = fopen(fname, "w");  
    
    if (!file) return 0;

    k = (1<<m) - 1 - 2*t;

    for (i = 0; i < data_length; i++) {
    
        M = allocate_vector(k);
        for (j = 0; j < k; j++) {
            M->data[j] = gen_symbol(m);
        }

        codeword = gen_codeword(m, t, add_one, M, g);
        
        R = add_error(e_prob, e_per_word, e_pmf, m, mode, codeword);
        
        for (j = 0; j < R->len; j++) {
            fprintf(file, "%i\n",R->data[j]);
        }

        free_vector(M);
        free_vector(codeword);
        free_vector(R);
    }

    free_vector(g);

    fclose(file);

    return 1;
    
}

